"""
    Title

    author: wxz
    date: 
    github: https://github.com/xinzwang
"""

import numpy as np
from time import *


class F_Rosenbrock(object):
    """
    Rosenbrock函数    用于测试最优化算法性能的非凸函数
    f_min = f(1,1) = 0
    """

    def __init__(self, n=10):
        self.n = n

    def forward1(self, x):
        out = 0
        for j in range(self.n - 1):
            out += 100 * (np.square(np.square(x[j]) - x[j + 1])) + np.square(x[j] - 1)
        return out

    def forward2(self, x):
        a = x[0:self.n - 1]
        b = x[1:self.n]
        return np.sum(100 * (np.square(np.square(a) - b)) + np.square(a - 1))

    def test(self):
        x = np.random.random([self.n])
        out1 = None
        out2 = None
        t1 = time()
        for i in range(100):
            out1 = self.forward1(x)
        t2 = time()
        for i in range(100):
            out2 = self.forward2(x)
        t3 = time()
        assert np.sum(out1 - out2) < 1e-9, print("Wrong!!!")

        print("Success!")
        print("forward1 runtime:", t2 - t1)
        print("forward2 runtime:", t3 - t2)
        print("acceleration ratio:", (t2 - t1) / (t3 - t2))
        return


f = F_Rosenbrock(n=1000)
f.test()

""" output
Success!
forward1 runtime: 0.5401351451873779
forward2 runtime: 0.000997781753540039
acceleration ratio: 541.3359617682198
"""
